import matplotlib.pyplot as plt
import cv2
import numpy as np
from pytorch_grad_cam import GradCAM
import torch

def create_mask(img1, img2):
    rows, cols, channels = img2.shape
    roi = img1[0:rows, 0:cols]
    img2gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(img2gray, 175, 255, cv2.THRESH_BINARY)
    mask_inv = cv2.bitwise_not(mask)
    img1_bg = cv2.bitwise_and(roi, roi, mask=mask)
    img2_fg = cv2.bitwise_and(img2, img2, mask=mask_inv)
    dst = cv2.add(img1_bg, img2_fg)
    img1[0:rows, 0:cols] = dst
    return img1


# gradcam预处理
def gradcam_preprocess(net, real_matrix):
    mx_tensor = torch.tensor(real_matrix, dtype=torch.float)  # 真实的散射矩阵
    cam = GradCAM(net, [net.layer4[-1]])
    img_tensor = mx_tensor.unsqueeze(0).unsqueeze(0)
    activation_map = cam(img_tensor)        # 输入到gradcam得到激活层
    plt.imsave('real_matrix.png', real_matrix, cmap='viridis')
    plt.imsave('activation_map.png', activation_map.squeeze(0), cmap='hot')
    real_matrix = plt.imread('real_matrix.png')
    ac_map = plt.imread('activation_map.png')
    visualization = cv2.addWeighted(real_matrix, 0.7, ac_map, 0.3, 0)
    return visualization


# 展示gradcam处理后的图像
def result_show(predic_matrix_pic,real_matrix_pic,visualization):

    plt.subplot(2,2,1)
    plt.imshow(predic_matrix_pic, cmap='viridis')     # 预测结果
    plt.title("predic_matrix_pic")
    plt.subplot(2,2,3)
    plt.imshow(real_matrix_pic, cmap='viridis')   # 真实结果
    plt.title("real_matrix_pic")
    plt.subplot(2,2,4)
    plt.imshow(visualization)
    plt.title("active_visual")      # 真实散射矩阵叠加激活层
    plt.colorbar()  # 添加颜色条
